import numpy as np
import matplotlib.pyplot as plt
from F4_性能指标转换 import F4_MPC_Matrices_PM
from F7_二次规划求解模块 import F7_MPC_Controller_withConstriants
from F6_约束条件矩阵转换模块 import F6_MPC_Matrices_Constraints
from scipy.interpolate import interp1d

plt.rcParams['font.sans-serif']=['SimHei'] #用来正常显示中文标签
plt.rcParams['axes.unicode_minus']=False #用来正常显示负号

# 系统定义
A = np.array([[-1,- 0.1, 0], [-0.5, 0, 0],[0, 0, -0.2]])
n = A.shape[0]
B = np.array([[1, 0], [0, 0],[0, 0.5]])
p = B.shape[1]
h_t_k = np.array([[0], [-2.0], [2.0]])
sigma_disturbance = 0.1  # 随机扰动的标准差

# 假定你的常量值如下
alpha = 0.3
rho = 2.2
h = 6
r = 6
mu = 0.9

# 权重设计
Q = 2 * np.eye(n)
S = 300 * np.eye(n)
R = 2 * np.eye(p)

# 系统初始化
x_0 = np.array([[1], [-1], [1]])
x = x_0

# 系统约束定义
x_low = np.array([-np.inf, -np.inf, 5]).reshape(-1, 1)
x_high = np.array([np.inf, np.inf, 20]).reshape(-1, 1)
u_low = np.array([-300, -50]).reshape(-1, 1)
u_high = np.array([500, 300]).reshape(-1, 1)

# 系统运行步数
k_steps =  50
x_history = np.zeros((n, k_steps))
x_history[:, 0] = x.reshape(-1)
u_history = np.zeros((p, k_steps))

# 预测区间
N_P = 3

# 计算二次规划需用到的矩阵
Phi, Gamma, Omega, Psi, M, H, F = F4_MPC_Matrices_PM(A,B,Q,R,S,N_P)

P, Beta_bar, b = F6_MPC_Matrices_Constraints(x_low,x_high,u_low,u_high,N_P,Phi,Gamma)

# 系统模拟
for k in range(k_steps-1):
    U, u = F7_MPC_Controller_withConstriants(x, M, H, F, P, Beta_bar, b, p, h_t_k)
    x = A @ x + B @ u + h_t_k

    # 计算y的值
    y = np.sqrt((alpha * x_history[2, :]) ** 2 + ((rho * h * x_history[1, :] ** 2 * r) / 2 * (1 - mu)) ** 2)
    if k == 20:
        # 更新h_t_k
        h_t_k[2,0] = 200 * np.random.normal(0, sigma_disturbance)  # 第三项为随机扰动
    else:
        h_t_k = np.array([[0], [-2.0], [2.0]]);

    x_history[:, k + 1] = x.reshape(-1)
    u_history[:,k] = u.reshape(-1)

#PID控制模拟

# 结果展示绘图
plt.figure(figsize=(5,5))
plt.plot(x_history[0,:], label='电机电流')
plt.plot(x_history[1,:], label='电机旋转角速度', linestyle='--')
plt.title('冠带层缠绕系统状态量变化图')
plt.plot(x_history[2,:], label='缠绕机头移动速度', linestyle='--')
plt.legend()
plt.grid(True)
plt.xlim([0, k_steps-1])
plt.ylim([-20, 20])
plt.tight_layout()
plt.show()

plt.step(range(k_steps), u_history[0,:], label='电机电压', where='post')
plt.title('冠带层缠绕系统输入量（电机电压）变化图')
plt.xlim([0, k_steps-1])
plt.ylim([-2, 2])
plt.tight_layout()
plt.show()

plt.step(range(k_steps), u_history[1,:], label='缠绕机头驱动力', where='post', linestyle='--')
plt.title('冠带层缠绕系统输入量（驱动力）变化图')
plt.legend()
plt.grid(True)
plt.xlim([0, k_steps-1])
plt.ylim([-20, 20])
plt.tight_layout()
plt.show()


y_interp = interp1d(range(k_steps), y, kind='cubic')  # cubic就是三次样条插值，你可以尝试其他的插值方法
new_steps = np.linspace(0, k_steps-1, k_steps*10)  # 创建新的横坐标，步长更小，点的数量更多，使得曲线更平滑
plt.figure(figsize=(5,5))
plt.plot(new_steps, y_interp(new_steps), label='实际系统的张力响应')  # 使用插值函数得到新的纵坐标，并画图
# 绘制设定值参考线18N
plt.axhline(y=18, color='r', linestyle='--', label='张力设定值18(N)')
# 绘制设定值参考线16N
plt.axhline(y=16, color='r', linestyle='--', label='张力设定值16(N)')
plt.legend()
plt.grid(True)
plt.xlim([0, k_steps-3])
plt.show()